import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from multilayer_perceptron import MultilayerPerceptron

data = pd.read_csv('mnist-demo.csv')

plt.figure(figsize=(20, 5))
pixel_datas = np.array(data.iloc[:100, 1:]).reshape(-1, 28, 28)
for i, pixel_data in enumerate(pixel_datas):
    plt.subplot(5, 20, i + 1)
    plt.imshow(pixel_data, cmap='gray')     # Grays
    plt.axis('off')
plt.show()

train_data = data.sample(frac=0.8)
test_data = data.drop(train_data.index)

X_train = train_data.values[:, 1:]
y_train = train_data.values[:, [0]]
X_test = test_data.values[:, 1:]
y_test = test_data.values[:, [0]]

layers = (784, 25, 10)
normalize_data = True
max_iterations = 1000
alpha = 0.1

multilayer_perceptron = MultilayerPerceptron(X_train, y_train, layers, normalize_data)
thetas, cost_history = multilayer_perceptron.train(max_iterations, alpha)
plt.plot(range(len(cost_history)), cost_history, 'b--')
plt.xlabel('Gradient Steps')
plt.ylabel('cost')
plt.show()

y_train_predicted = multilayer_perceptron.predict(X_train)
y_test_predicted = multilayer_perceptron.predict(X_test)

train_acc = np.sum(y_train_predicted == y_train) / y_train.shape[0] * 100
test_acc = np.sum(y_test_predicted == y_test) / y_test.shape[0] * 100
print('训练集准确率 :', train_acc, '%')
print('测试集准确率 :', test_acc, '%')

plt.figure(figsize=(10, 10))
pixel_datas = np.array(X_test[:36]).reshape(-1, 28, 28)
for i, pixel_data in enumerate(pixel_datas):
    plt.subplot(6, 6, i + 1)
    color_map = 'Greens' if y_test_predicted[i] == y_test[i] else 'Reds'
    plt.imshow(pixel_data, cmap=color_map)
    plt.title(f'{y_test_predicted[i]}')
    plt.axis('off')
plt.show()
